import pandas as pd
import geopandas as gpd
import folium
import compare_str as cp

# 读取Excel文件中的数据
file_path = 'E:\\大二上\\Weiguan\\data\\population_of_Beijing.xlsx'
df = pd.read_excel(file_path)

# 只处理2020年的人口数据
df_2020 = df[['Status', 'Native', 'Population Census 2020-11-01']]

# 初始化一个空字典来存储每个区的人口数据
district_data = {}

# 遍历数据框的每一行
for index, row in df_2020.iterrows():
    district = row['Native']
    if district not in district_data and (row['Status'] != 'District'):
        district_data[district] = []
    district_data[district].append((row['Native'], row['Population Census 2020-11-01']))

# 读取北京边界的shp文件
shp_path = "E:\\大二上\\Weiguan\\data\\北京市乡镇街区边界\\BJ_country.shp"
gdf = gpd.read_file(shp_path)

print(gdf.columns)
print('\n')
print(gdf['NAME'])
print('\n')
print(district_data.items())

# 创建一个以北京市为中心的地图
beijing_map = folium.Map(location=[39.9042, 116.4074], zoom_start=10)

# 将北京边界添加到地图中
folium.GeoJson(gdf.to_json(), name="Beijing Boundaries").add_to(beijing_map)

# 将每个区的人口数据添加到地图上
for district, populations in district_data.items():
    for name in gdf['NAME']:
        if cp.are_strings_equal_ignore_whitespace(district, name):
            boundary = gdf[gdf['NAME'] == name]
            if not boundary.empty:
                for _, population in populations:
                    folium.GeoJson(
                        data=boundary.to_json(),
                        style_function=lambda x: {'fillColor': 'blue', 'color': 'blue', 'weight': 1, 'fillOpacity': 0.5},
                        tooltip=folium.Tooltip(f"{district}: {population}人")
                    ).add_to(beijing_map)

# 保存地图到HTML文件
beijing_map.save("beijing_population_distribution_map.html")